###在Kihwah's lab，目前所有需要awg的实验都使用本文件中的硬件接口，所以将本文件放在cion目录下
import types
import ctypes
import importlib
import sys
import numpy as np

# tensorflow的lazyloader用法：lazyloader一个module时，只有真正用到该module时才会load
class LazyLoader(types.ModuleType):
    """Lazily import a module, mainly to avoid pulling in large dependencies.
    ​`contrib`, and `ffmpeg` are examples of modules that are large and not alwaysneeded, and this allows them to only be loaded when they are used."""
    def __init__(self, local_name, parent_module_globals, name):
        self._local_name = local_name
        self._parent_module_globals = parent_module_globals
        super(LazyLoader, self).__init__(name)

    def _load(self):
        # Import the target module and insert it into the parent's namespace
        module = importlib.import_module(self.__name__)
        self._parent_module_globals[self._local_name] = module
        # Update this object's dict so that if someone keeps a reference to the
        # LazyLoader, lookups are efficient (__getattr__ is only called on lookups
        # that fail).
        self.__dict__.update(module.__dict__)
        return module

    def __getattr__(self, item):
        module = self._load()
        return getattr(module, item)

    def __dir__(self):
        module = self._load()
        return dir(module)


#这里我们使用的是Spectrum Instrumentation公司的任意波形发生器（AWG）
#对于使用其他种类awg的用户，请自行定义好相关接口，即重构下述的AWG_Control类
class AWG_Control:
    #初始化，通过ip地址连接到AWG硬件
    def __init__(self, ip_address):
        spctr = LazyLoader("spctr", globals(), "cion.hardware.awg.pyspcm")
        spctr_tools = LazyLoader("spctr_tool", globals(), "cion.hardware.awg.spcm_tools")
        # ip_address = '192.168.2.30'
        address_string = 'TCPIP::'+ip_address
        self.hCard = spctr.spcm_hOpen (spctr.create_string_buffer (bytes(address_string, encoding='gbk')))
        self.channel_list = [spctr.CHANNEL0, spctr.CHANNEL1, spctr.CHANNEL2, spctr.CHANNEL3, spctr.CHANNEL4, spctr.CHANNEL5, \
            spctr.CHANNEL6, spctr.CHANNEL7, spctr.CHANNEL8, spctr.CHANNEL9, spctr.CHANNEL10, spctr.CHANNEL11, spctr.CHANNEL12, spctr.CHANNEL13,\
            spctr.CHANNEL14, spctr.CHANNEL15, spctr.CHANNEL16, spctr.CHANNEL17, spctr.CHANNEL18, spctr.CHANNEL19, spctr.CHANNEL20, spctr.CHANNEL21,\
            spctr.CHANNEL22, spctr.CHANNEL23, spctr.CHANNEL24, spctr.CHANNEL25, spctr.CHANNEL26, spctr.CHANNEL27, spctr.CHANNEL28, spctr.CHANNEL29,\
            spctr.CHANNEL30, spctr.CHANNEL31]

        self.qwChEnable = spctr.CHANNEL0
        self.channel = [0]

        #print type and serial number of AWG
        lCardType = spctr.int32 (0)
        spctr.spcm_dwGetParam_i32 (self.hCard, spctr.SPC_PCITYP, spctr.byref (lCardType))
        lSerialNumber = spctr.int32 (0)
        spctr.spcm_dwGetParam_i32 (self.hCard, spctr.SPC_PCISERIALNO, spctr.byref (lSerialNumber))
        lFncType = spctr.int32 (0)
        spctr.spcm_dwGetParam_i32 (self.hCard, spctr.SPC_FNCTYPE, spctr.byref (lFncType))
        sCardName = spctr_tools.szTypeToName (lCardType.value)
        sys.stdout.write("Found: {0} sn {1:05d}\n".format(sCardName,lSerialNumber.value))

        # set samplerate to 625 MHz, no clock output
        spctr.spcm_dwSetParam_i64 (self.hCard, spctr.SPC_SAMPLERATE, spctr.MEGA(625))
        spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_CLOCKOUT,   0)

    # 关闭这个awg接口对象
    def close(self):
        spctr = LazyLoader("spctr", globals(), "cion.hardware.awg.pyspcm")
        spctr.spcm_vClose(self.hCard)

    def error_check(self):
        szErrorTextBuffer = spctr.create_string_buffer (spctr.ERRORTEXTLEN)
        dwError = spctr.spcm_dwGetErrorInfo_i32 (self.hCard, None, None, szErrorTextBuffer)
        assert dwError == 0, f"An error happend! Error code: {dwError}"
    # 激活awg板上对应的channels
    def enable_channels(self, channels=[0],soft_flag=False):
        assert type(channels) == list and len(channels) <= 4
        self.channel = channels
        register_list = [spctr.SPC_ENABLEOUT0, spctr.SPC_ENABLEOUT1, spctr.SPC_ENABLEOUT2, spctr.SPC_ENABLEOUT3]
        self.qwChEnable = self.channel_list[channels[0]] # channels 0,2,3 -> 1101 = 13
        channels_remained = [0,1,2,3]
        [channels_remained.remove(i) for i in channels]
        ### set the turned-on channels, where 'qwChEnable = 15 = 1111' means opening all 4 channels
        self.qwChEnable = 0
        for i in channels:
            self.qwChEnable = self.qwChEnable | self.channel_list[i]
            
        dwError = spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_M2CMD, spctr.M2CMD_CARD_STOP)
        spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_CARDMODE,    spctr.SPC_REP_STD_SINGLERESTART) #SPC_REP_STD_SINGLERESTART 
        # set samplerate to 625 MHz, no clock output
        spctr.spcm_dwSetParam_i64 (self.hCard, spctr.SPC_SAMPLERATE, spctr.MEGA(625))
        spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_CLOCKOUT,   0)
        spctr.spcm_dwSetParam_i64 (self.hCard, spctr.SPC_CHENABLE,    self.qwChEnable)

        register_list_enable = [register_list[i] for i in channels]
        for register in register_list_enable:
            spctr.spcm_dwSetParam_i64 (self.hCard, register,  1)

        register_list_disenable = [register_list[i] for i in channels_remained]
        for register in register_list_disenable:
            spctr.spcm_dwSetParam_i64 (self.hCard, register,  0)

        if soft_flag==True:
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_ORMASK,      spctr.SPC_TMASK_SOFTWARE) #SPC_TMASK_SOFTWARE表示软件trigger，在card运行后会立刻给一个trigger
            #spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_ORMASK,      spctr.SPC_TMASK_EXT0) #在trigger0收到信号后触发
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_ANDMASK,     0)
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_CH_ORMASK0,  0)
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_CH_ORMASK1,  0)
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_CH_ANDMASK0, 0)
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_CH_ANDMASK1, 0)
        else:
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_ORMASK,      spctr.SPC_TMASK_EXT0)         # external trigger
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_EXT0_MODE,   spctr.SPC_TM_POS)             # trigger on positive edge
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_EXT0_LEVEL0, 1800)                   # lower trigger at 1500mV
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_EXT0_LEVEL1, 2000)                   # upper trigger at 1800mV
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIG_ANDMASK,     0)  

        lSetChannels = spctr.int32 (0)
        spctr.spcm_dwGetParam_i32 (self.hCard, spctr.SPC_CHCOUNT,     spctr.byref (lSetChannels))
        print(lSetChannels.value, "channels are enabled.")
        error_msg = self.error_check()
        assert error_msg==None, error_msg

    #将ndarray转换成awg dll接口能识别的数据类型
    def to_DAC_readable(self, x):
        spctr = LazyLoader("spctr", globals(), "cion.hardware.awg.pyspcm")
        length = 32*len(self.channel)
        x = np.transpose(x)
        x = x.flatten()
        x = np.pad(x,(0,length-len(x)%length))
        return np.floor(32767*x).astype(spctr.int16)

    # directly send multi-channel data to awg board
    # n*m ndarray for n chaneels
    # 直接发送多个channels的数据到awg板
    # data_array的size是n*m：n个channel，每个channel有m个数据点
    def send_data(self, data_array, loop=1, print_flag=True):
        spctr = LazyLoader("spctr", globals(), "cion.hardware.awg.pyspcm")
        spctr_tools = LazyLoader("spctr_tool", globals(), "cion.hardware.awg.spcm_tools")
        data_array = self.to_DAC_readable(data_array)
        assert type(data_array[0]) == np.int16
        #保证总数据点的个数是32的整数倍
        assert len(data_array) % 32 == 0
        # setup the mode
        dwError = spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_M2CMD, spctr.M2CMD_CARD_STOP) #关闭card

        llMemSamples = spctr.int64 (int(len(data_array)/len(self.channel))) # Size of memory must a multiple of 32
        llLoops = spctr.int64 (loop) # loop only once #如果想要连续输出就改成 0
        spctr.spcm_dwSetParam_i64 (self.hCard, spctr.SPC_MEMSIZE,     llMemSamples)
        spctr.spcm_dwSetParam_i64 (self.hCard, spctr.SPC_LOOPS,       llLoops)
        lSetChannels = spctr.int32 (0)
        spctr.spcm_dwGetParam_i32 (self.hCard, spctr.SPC_CHCOUNT,     spctr.byref (lSetChannels))
        lBytesPerSample = spctr.int32 (0)
        spctr.spcm_dwGetParam_i32 (self.hCard, spctr.SPC_MIINST_BYTESPERSAMPLE,  spctr.byref (lBytesPerSample))
        spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TRIGGEROUT,       0)

        lChannel = spctr.int32 (0)
        spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_AMP0 + lChannel.value * (spctr.SPC_AMP1 - spctr.SPC_AMP0),\
                                   spctr.int32 (1000))

        # we define the buffer for transfer and start the DMA transfer
        if print_flag==True:
            # setup software buffer
            qwBufferSize = spctr.uint64 (llMemSamples.value * lBytesPerSample.value * lSetChannels.value)
            # get pointer of input data
            pvBuffer = ctypes.c_void_p (data_array.__array_interface__['data'][0])
            sys.stdout.write("Starting the DMA transfer and waiting until data is in board memory\n")
            spctr.spcm_dwDefTransfer_i64 (self.hCard, spctr.SPCM_BUF_DATA, spctr.SPCM_DIR_PCTOCARD, spctr.int32 (0), pvBuffer, spctr.uint64 (0), qwBufferSize)
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_M2CMD, spctr.M2CMD_DATA_STARTDMA | spctr.M2CMD_DATA_WAITDMA)
            sys.stdout.write("... data has been transferred to board memory\n")

            # We'll start and wait until the card has finished or until a timeout occurs
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TIMEOUT, 10000)
            sys.stdout.write("\nStarting the card and waiting for ready interrupt\n(continuous and single restart will have timeout)\n")
            dwError = spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_M2CMD, spctr.M2CMD_CARD_START \
                                             | spctr.M2CMD_CARD_ENABLETRIGGER | spctr.M2CMD_CARD_WAITREADY)
        else:
            # setup software buffer
            qwBufferSize = spctr.uint64 (llMemSamples.value * lBytesPerSample.value * lSetChannels.value)
            # get pointer of input data
            pvBuffer = ctypes.c_void_p (data_array.__array_interface__['data'][0])
            spctr.spcm_dwDefTransfer_i64 (self.hCard, spctr.SPCM_BUF_DATA, spctr.SPCM_DIR_PCTOCARD, spctr.int32 (0), pvBuffer, spctr.uint64 (0), qwBufferSize)
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_M2CMD, spctr.M2CMD_DATA_STARTDMA | spctr.M2CMD_DATA_WAITDMA)

            # We'll start and wait until the card has finished or until a timeout occurs
            spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_TIMEOUT, 10000)
            dwError = spctr.spcm_dwSetParam_i32 (self.hCard, spctr.SPC_M2CMD, spctr.M2CMD_CARD_START \
                                             | spctr.M2CMD_CARD_ENABLETRIGGER | spctr.M2CMD_CARD_WAITREADY)
            
            #M2CMD_CARD_WAITTRIGGER
            

        error_msg = self.error_check()
        assert error_msg==None, error_msg
